package com.gmrz.webauthn.protocol.v1.processor.attestation;

import com.gmrz.uaf.protocol.v1.json.UAFSchemaBuilder;
import com.gmrz.webauthn.common.WebAuthnConstants;
import com.gmrz.webauthn.protocol.v1.json.WebAuthnSchemaBuilder;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnErrorCode;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnException;
import com.gmrz.webauthn.protocol.v1.schema.*;
import com.google.gson.Gson;
import com.google.inject.Inject;
import com.gmrz.util.Convert;
import com.gmrz.util.Streams;
import com.gmrz.uaf.crypto.CryptoEngine;
import com.gmrz.uaf.crypto.spi.handlers.AttestationCertHandler;
import com.gmrz.webauthn.common.WebAuthnRuntimeException;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.GeneralSecurityException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1OctetString;
import org.bouncycastle.asn1.x509.Extension;
import org.bouncycastle.cert.X509CertificateHolder;

public class WebAuthnBaseAssertionProcessor
{
    private static final Logger LOG = LogManager.getLogger(WebAuthnBaseAssertionProcessor.class);
    protected CryptoEngine cryptoEngine = null;
    protected WebAuthnConstants.AttestationResult attResult = WebAuthnConstants.AttestationResult.SUCCESS;

    @Inject
    public void setCryptoEngine(CryptoEngine cryptoEngine)
    {
        this.cryptoEngine = cryptoEngine;
    }

    protected byte[] getClientDataHash(byte[] clientDataBytes, HashAlgorithmIdentifier algorithm)
            throws WebAuthnRuntimeException
    {
        byte[] clientDataHash = null;
        String errorMsg = null;
        try
        {
            switch (algorithm)
            {
                case S256:
                    clientDataHash = this.cryptoEngine.getDataHashHandler().hash(clientDataBytes);

                    break;
                default:
                    errorMsg = "Failed to calculate the client data hash. Algorithm" + algorithm + "is not supported";

                    LOG.error(errorMsg);
                    throw new WebAuthnRuntimeException(WebAuthnErrorCode.STATUS_CLIENTDATA_HASHING_FAILED, errorMsg);
            }
            return clientDataHash;
        }
        catch (GeneralSecurityException e)
        {
            errorMsg = "Exception is thrown when hashing client data: " + e.getMessage();

            LOG.error(errorMsg);
            throw new WebAuthnRuntimeException(WebAuthnErrorCode.STATUS_CLIENTDATA_HASHING_FAILED, errorMsg);
        }
    }

    public void validateCertificate(List<Certificate> chain, WebAuthnMetadata webAuthnMetadata, boolean policyQualifiersRejected)
            throws WebAuthnException
    {
        if ((chain == null) || (chain.isEmpty()))
        {
            String errorMsg = "Could not find attestation certificate.";
            logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
        }
        String normalizeAAGUID =  webAuthnMetadata.getAaguiid();
        Gson gson = UAFSchemaBuilder.getGson();
        WebAuthenticatorSpec spec = gson.fromJson(webAuthnMetadata.getJsonData(),WebAuthenticatorSpec.class);
        if ((spec == null) || (spec.getRootCertificates() == null) || (spec.getRootCertificates().isEmpty()))
        {
            String errorMsg = "Trust Anchors not available for validation";
            logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_TRUST_ANCHOR_NOT_FOUND);
        }
        else
        {
            try
            {
                validateCertChain(chain, spec.getRootCertificates(), normalizeAAGUID, policyQualifiersRejected);
            }
            catch (WebAuthnException e)
            {
                String errorMsg = "Failed to validate ceritifcate chain against trust anchors specified in metadata";
                logAndThrow(errorMsg, (WebAuthnErrorCode)e.getErrorCode(), e);
            }
        }
    }

    private void validateCertChain(List<Certificate> chain, List<Certificate> rootCerts, String aaguid, boolean policyQualifiersRejected)
            throws WebAuthnException
    {
        boolean result = false;
        String errormsg = "Failed to validate attestation certificate- ";
        try
        {
            AttestationCertHandler handler = this.cryptoEngine.getAttestationCertHandler();
            if (handler != null) {
                result = handler.validateCertificate(chain, rootCerts, aaguid);
            }
        }
        catch (RuntimeException cpve)
        {
            errormsg = errormsg + cpve.getMessage();
            int index = 0;
            if ((index >= 0) && (index < chain.size())) {
                LOG.error("Certificate at index[" + index + "] in the chain is not good.", ((Certificate)chain.get(index)).toString());
            }
            logAndThrow(errormsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED, cpve);
        }
        catch (Exception gse)
        {
            errormsg = errormsg + gse.getMessage();
            logAndThrow(errormsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED, gse);
        }
        if (!result)
        {
            LOG.error("Failed to validate attestation certificate for AAGUID[{}]", aaguid);
            throw new WebAuthnException(WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
        }
    }

    public void validateSignature(int signAlgo, Certificate attestationCert, byte[] signedData, byte[] signature)
            throws WebAuthnException
    {
        if (attestationCert == null)
        {
            String errormsg = "Couldn't find attestation cert for Basic Full Attestation.";
            logAndThrow(errormsg, WebAuthnErrorCode.ERROR_ATTESTATION_SIGNATURE_VALIDATION_FAILED);
        }
        boolean result = false;
        try
        {
            result = this.cryptoEngine.verifySignature(signAlgo, attestationCert, signedData, signature);
        }
        catch (Exception cpve)
        {
            logAndThrow(cpve.getMessage(), WebAuthnErrorCode.ERROR_ATTESTATION_SIGNATURE_VALIDATION_FAILED, cpve);
        }
        if (!result) {
            throw new WebAuthnException(WebAuthnErrorCode.ERROR_ATTESTATION_SIGNATURE_VALIDATION_FAILED);
        }
    }

    protected void validateSignature(int keyFormat, byte[] pubKey, int signAlgo, byte[] signedData, byte[] signature)
            throws WebAuthnException
    {
        boolean result = false;
        try
        {
            result = this.cryptoEngine.verifySignature(keyFormat, pubKey, signAlgo, signedData, signature);
        }
        catch (Exception cpve)
        {
            logAndThrow(cpve.getMessage(), WebAuthnErrorCode.ERROR_SIGNATURE_VALIDATION_FAILED, cpve);
        }
        if (!result) {
            throw new WebAuthnException(WebAuthnErrorCode.ERROR_SIGNATURE_VALIDATION_FAILED);
        }
    }

    protected void logAndUpdateResult(String msg, WebAuthnErrorCode error, Exception e)
            throws WebAuthnException
    {
        LOG.error(msg, e);
        this.attResult = WebAuthnConstants.AttestationResult.FAILED;
    }

    protected void logAndThrow(String errorMsg, WebAuthnErrorCode error, Exception e)
            throws WebAuthnException
    {
        LOG.error(errorMsg, e);
        throw new WebAuthnException(error, errorMsg, e);
    }

    protected void logAndThrow(String errorMsg, WebAuthnErrorCode error)
            throws WebAuthnException
    {
        LOG.error(errorMsg);
        throw new WebAuthnException(error, errorMsg);
    }

    protected void logAndUpdateResult(String msg, WebAuthnErrorCode error)
            throws WebAuthnException
    {
        LOG.error(msg);
        this.attResult = WebAuthnConstants.AttestationResult.FAILED;
    }

    protected List<Certificate> getDecodedX509Certs(List<Object> x509Certs)
            throws WebAuthnException
    {
        List<Certificate> certsList = new ArrayList();
        AttestationCertHandler attestationCertHandler;
        String errorMsg;
        try
        {
            attestationCertHandler = this.cryptoEngine.getAttestationCertHandler();
            if (attestationCertHandler == null)
            {
                errorMsg = "Certificate decoding failed due to null AttestationCertHandler.";
                logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED);
            }
            else
            {
                for (Object encodedBase64Cert : x509Certs)
                {
                    Certificate cert = null;
                    byte[] encodedBase64CertBytes = (byte[])encodedBase64Cert;
                    cert = attestationCertHandler.decode(encodedBase64CertBytes);
                    certsList.add(cert);
                }
            }
        }
        catch (CertificateException ex)
        {
            errorMsg = "Exception while reading the certificates from the extension value";
            LOG.error(errorMsg, ex);
            throw new WebAuthnException(WebAuthnErrorCode.ERROR_ATTESTATION_MALFORMED_STATEMENT, errorMsg);
        }
        return certsList;
    }

    protected void verifyAAGUIDInCertExtentsion(X509Certificate x509Cert, CredentialAttestation credentialAttestation)
            throws WebAuthnException
    {
        if (x509Cert == null)
        {
            LOG.error("Null certificate found, can't perform verification");
            return;
        }
        try
        {
            LOG.info("Verifying AAGUID of extension with AAGUID of authenticator data");

            ASN1ObjectIdentifier oid = new ASN1ObjectIdentifier("1.3.6.1.4.1.45724.1.1.4");
            Extension ext = new X509CertificateHolder(x509Cert.getEncoded()).getExtension(oid);
            if (ext == null)
            {
                LOG.debug("Verification of AAGUID extension is not required as optional id-fido-gen-ce-aaguid extention is not present in X509 attestation certificate.");

                return;
            }
            if (ext.isCritical())
            {
                String errorMsg = "id-fido-gen-ce-aaguid extention is marked as critical, must be marked non critical";
                LOG.error(errorMsg);
            }
            ASN1OctetString value = ext.getExtnValue();
            byte[] octetValue = value.getOctets();

            LOG.debug("AAGUID extension value: {}", Convert.toHex(octetValue));
            if (octetValue.length < 18)
            {
                String errorMsg = "AAGUID extension data is not as expected, doesn't contains the minimum required bytes (i.e 18 containig OID sequence and 16 bytes AAGUID double octet:";

                LOG.error(errorMsg);
            }
            else
            {
                AAGUID aaguidFromExtension = new AAGUID(fetchAAGUIDFromOctet(octetValue));

                AAGUID authenticaorAAGUID = credentialAttestation.getAuthenticatorData().getAttestation().getAaguid();

                LOG.info("Hex format AAGUID from Extention is: " + aaguidFromExtension + ", AAGUID from Authenticator data is: " + authenticaorAAGUID);
                if (!authenticaorAAGUID.equals(aaguidFromExtension))
                {
                    String errorMsg = "AAGUID from ext data does not match Authenticator AAGUID";
                    LOG.error(errorMsg);
                }
            }
        }
        catch (Exception e)
        {
            String errorMsg = "Failed to fetch AAGUID from ext data";
            logAndThrow(errorMsg, WebAuthnErrorCode.ERROR_ATTESTATION_CERT_VALIDATION_FAILED, e);
        }
    }

    private byte[] fetchAAGUIDFromOctet(byte[] octetValue)
    {
        if (octetValue.length < 18) {
            throw new IllegalArgumentException("Must provide an OCTET STRING (DER encoded format) byte array of at least 35 length.");
        }
        byte[] aaguid = Arrays.copyOfRange(octetValue, 2, 18);

        return aaguid;
    }

    protected int getSignAlgFromAttestaionMap(Map<String, Object> attStmt)
    {
        int coseAlgIdentifier =
                ((BigInteger) attStmt.get("alg")).intValue();
        LOG.debug("coseAlgIdentifier from Attestation statement:" + coseAlgIdentifier);

        int signAlgo = SignAlgorithmIdentifier.forCoseAlgID(coseAlgIdentifier).getSignAlgID();

        return signAlgo;
    }

    protected byte[] getSignedData(byte[] rawAuthData,byte[] clientData)
            throws IOException, NoSuchAlgorithmException, WebAuthnException
    {
        ByteArrayOutputStream outputStream = null;
        try
        {
            CollectedClientData collectedClientData = WebAuthnSchemaBuilder.getGson().fromJson(new String(clientData), CollectedClientData.class);

            outputStream = new ByteArrayOutputStream();

            outputStream.write(rawAuthData);

            byte[] clientDataHash = getClientDataHash(clientData, collectedClientData
                    .getHashAlg());
            LOG.debug(" client data hash size : {}", Integer.valueOf(clientDataHash.length));

            outputStream.write(clientDataHash);
            byte[] generatedSignData = outputStream.toByteArray();

            return generatedSignData;
        }
        finally
        {
            Streams.close(outputStream);
        }
    }
}

